#region File Description
//-----------------------------------------------------------------------------
// AnimationPlayer.cs
//
// Microsoft XNA Community Game Platform
// Copyright (C) Microsoft Corporation. All rights reserved.
//-----------------------------------------------------------------------------
#endregion

#region Using Statements
using System;
using System.Collections.Generic;
using Microsoft.Xna.Framework;
#endregion


namespace SkinnedModel
{

    /// <summary>
    /// The IK Solver is in charge of goal-directed animation
    ///   internally, goals manipulate a skeleton (of quaternions)
    ///   the quaternions are turned into joint matrices (the final vertex shader understands matrices, not quat)
    /// </summary>
    public class IKSolver
    {
        #region Fields

        // Information about the currently playing animation clip.
        TimeSpan currentTimeValue;

        // Current animation transform matrices.
        Matrix[] boneTransforms;  //this is a bit hacky - bone transforms is sticking around so that 
        // we can do quick "set from animation data" operations

        //these are for setting final bone transforms
        Matrix[] worldTransforms;
        Matrix[] skinTransforms;

        // Backlink to the bind pose and skeleton hierarchy data.
        SkinningData skinningDataValue;

        //These are the "in use" joints
        //List<int> controlledJoints = new List<int>();

        //IK goal list
        List<IKGoal> goals;
        
        //skeleton - a quaternion based representation of our articulated character
        Skeleton skeleton;
        Dictionary<int, List<int>> children; //reverse of skeleton parent info
        Matrix rootTransform = new Matrix();

        #endregion


        /// <summary>
        /// Constructs a new animation player.
        /// </summary>
        public IKSolver(SkinningData skinningData)
        {
            if (skinningData == null)
                throw new ArgumentNullException("skinningData");

            skinningDataValue = skinningData;

            boneTransforms = new Matrix[skinningData.BindPose.Count];
            worldTransforms = new Matrix[skinningData.BindPose.Count];
            skinTransforms = new Matrix[skinningData.BindPose.Count];

            skinningData.BindPose.CopyTo(boneTransforms, 0);
            goals = new List<IKGoal>();

            //TODO: read in connectivity
            int[] parents = new int[skinningData.SkeletonHierarchy.Count];

            skinningData.SkeletonHierarchy.CopyTo(parents, 0);

            skeleton = new Skeleton(boneTransforms, parents, skinningData.BoneNames);

            children = new Dictionary<int, List<int>>();
            for (int i = 0; i < parents.Length; i++)
            {
                children[i] = new List<int>();
            }

            for (int i = 0; i < parents.Length; i++)
            {
                int parent = parents[i];
                if(parent > -1)
                    children[parent].Add(i);
            }

        }

        protected List<int> GetChildren(int parentIdx)
        {
            return children[parentIdx];
        }

        public List<int> GetAllDescentdents(string root)
        {
            int rootIdx = skeleton.boneNames[root];
            List<int> retList = new List<int>();
            List<int> additionalCtl = new List<int>();
            //process controlled joints
            retList = children[rootIdx];

            do
            {
                additionalCtl.Clear();
                foreach (int c in retList)
                {
                    List<int> chList = children[c];
                    foreach (int ch in chList)
                    {
                        if (!retList.Contains(ch))
                        {
                            additionalCtl.Add(ch);
                        }
                    }
                }
                foreach (int c in additionalCtl)
                {
                    if (!retList.Contains(c))
                    {
                        retList.Add(c);
                    }
                }
            } while (additionalCtl.Count > 0);

            return retList;
        }

        //DEBUG for getting a the most recent rotation axis
        //
        public Vector3 _DBG_GetAxis()
        {
            return skeleton.axis;
        }

        //reverse lookup to get a parent joint name from a child name
        public string GetParentName(string effectorName)
        {
            int effectorId = -1;
            skeleton.boneNames.TryGetValue(effectorName, out effectorId);
            int jointId = skeleton.parents[effectorId];

            //reverse lookup
            foreach (KeyValuePair<string, int> pair in skeleton.boneNames)
            {
                if (pair.Value == jointId)
                    return pair.Key;
            }

            return "";
        }

        /*
        public List<int> GetControlledJoints()
        {
            
            return controlledJoints;
        }
         */

        ///<summary>
        ///Gets the aim of an effector
        /// aim = normalize(effector-parent)        
        /// </summary>
        /// <param name="effectorName"></param>
        /// <param name="aim"></param>
 
        public void GetEffectorAim(string effectorName, out Vector3 aim)
        {
            int effectorId = -1;
            skeleton.boneNames.TryGetValue(effectorName, out effectorId);            
            int jointId = skeleton.parents[effectorId];

            skeleton.GetEffectorAim(effectorId, jointId, out aim);

        }


        //gets an effector joint
        //also gets parent
        public void GetEffectorPosition(string effectorName, out Vector3 parent, out Vector3 effector)
        {
            
            int effectorId = -1;
            skeleton.boneNames.TryGetValue(effectorName, out effectorId);
            int parentId = skeleton.parents[effectorId];

            //TODO: combine these into something more efficient
            skeleton.GetEffectorPosition(parentId, out parent);

            skeleton.GetEffectorPosition(effectorId, out effector);

        }

        //gets an effector joint
        //also gets parent
        public void GetEffectorPosition(string effectorName, out Vector3 effector)
        {

            int effectorId = -1;
            skeleton.boneNames.TryGetValue(effectorName, out effectorId);

            skeleton.GetEffectorPosition(effectorId, out effector);

        }

        //gets an effector joint
        //also gets parent
        public void GetEffectorPosition(int effectorId, out Vector3 effector)
        {
            skeleton.GetEffectorPosition(effectorId, out effector);

        }


        public int GetNumJoints()
        {
            return skeleton.joints.Length;
        }

        public void GetEffectorRotation(int jointID, out Quaternion q)
        {
            skeleton.GetJointRotation(jointID, out q);
        }

        /// <summary>
        ///  Sets character to first key frame of an animation clip
        /// Useful for posing
        /// </summary>
        public void SetPoseFromFirstKey(AnimationClip clip)
        {
            if (clip == null)
                throw new ArgumentNullException("clip");


            currentTimeValue = TimeSpan.Zero;
            int currentKeyframe = 0;

            // Initialize bone transforms to the bind pose.
            skinningDataValue.BindPose.CopyTo(boneTransforms, 0);

            // Read keyframe matrices.
            IList<Keyframe> keyframes = clip.Keyframes;

            while (currentKeyframe < keyframes.Count)
            {
                Keyframe keyframe = keyframes[currentKeyframe];

                // Stop when we've read up to the current time position.
                if (keyframe.Time > TimeSpan.Zero)
                    break;

                // Use this keyframe.
                boneTransforms[keyframe.Bone] = keyframe.Transform;

                currentKeyframe++;
            }
            skeleton = new Skeleton(boneTransforms, skeleton.parents, skeleton.boneNames);
        }

        /// <summary>
        /// Adds a One Joint Goal (currently aim based)
        /// </summary>
        /// <param name="effector"></param>
        /// <param name="aim"></param>
        public void AddOneJointAimGoal(string effector, Vector3 aim)
        {
            int effectorIdx = -1;

            skeleton.boneNames.TryGetValue(effector, out effectorIdx);

            Goal1JointAim G = new Goal1JointAim(effectorIdx, aim);
            AddGoal(G);
        }

        public void AddReachGoal(string effector, string root, Vector3 target)
        {

            int effectorIdx = -1, rootIdx = -1;

            skeleton.boneNames.TryGetValue(effector, out effectorIdx);
            skeleton.boneNames.TryGetValue(root, out rootIdx);

            ReachGoal rg = new ReachGoal(effectorIdx, rootIdx, target);
            AddGoal(rg);

        }

        /// <summary>
        /// Sets an absolute orientation goal 
        /// Basically, joint -> assume this orientation
        /// </summary>
        /// <param name="effector"></param>
        /// <param name="ornt"></param>
        public void AddAbsoluteGoal(string effector, ref Quaternion ornt)
        {
            int effectorIdx = -1;

            skeleton.boneNames.TryGetValue(effector, out effectorIdx);

            AbsoluteOrientationGoal G = new AbsoluteOrientationGoal(effectorIdx, ornt);
            AddGoal(G);
        }

        /// <summary>
        /// Tells solver to apply a relative rotation to a joint
        /// </summary>
        /// <param name="effector"></param>
        /// <param name="q"></param>

        public void AddRelativeGoal(string effector, ref Quaternion q)
        {
            int effectorIdx = -1;
            skeleton.boneNames.TryGetValue(effector, out effectorIdx);

            RelativeOrientationGoal G = new RelativeOrientationGoal(effectorIdx, q);
            AddGoal(G);
        }

        /// <summary>
        /// Advances the current animation position.
        /// </summary>
        public void Update(TimeSpan time,
                           ref Matrix rootTransform)
        {
            UpdateIK(ref rootTransform);
            UpdateWorldTransforms();
            UpdateSkinTransforms();
        }

        protected void AddGoal(IKGoal g)
        {
            goals.Add(g);
        }

        protected void DeleteGoal(IKGoal g)
        {
            goals.Remove(g);
        }

        public void UpdateMatrices(ref Matrix[] animMatrices, ref float[] jointBlends)
        {
            Quaternion q = new Quaternion();
            Vector3 s = new Vector3();
            Vector3 t = new Vector3();

            for (int i = 0; i < jointBlends.Length; i++)
            {
                if (jointBlends[i] == 0.0f) //Are we under Anim Control?
                {
                    animMatrices[i].Decompose(out s, out q, out t);
                    boneTransforms[i] = animMatrices[i];
                    q.Normalize();

                    Matrix DBG = Matrix.CreateFromQuaternion(q);

                    Vector3 v = new Vector3();
                    TMath.Quaternion2Eulers(ref q, out v);
                    v *= 180.0f / MathHelper.Pi;
                    //Console.WriteLine("Initial rotations"); Console.WriteLine(v);

                    skeleton.joints[i] = new Joint(q, t);
                }
            }
            
        }

        /// <summary>
        /// Solve each IK goal sequentially
        /// Then update bone transforms
        /// TODO: update only tagged joints
        /// </summary>
        public void UpdateIK(ref Matrix rootTransform)
        {
            this.rootTransform = rootTransform;
            //controlledJoints.Clear();

            foreach (IKGoal g in goals)
            {
                g.Solve(ref skeleton);
                //g.GetControlledJoints(ref controlledJoints, ref skeleton);
            }
            goals.Clear();

            /*
            List<int> additionalCtl = new List<int>();
            //process controlled joints
            do
            {
                additionalCtl.Clear();
                foreach (int c in controlledJoints)
                {
                    List<int> chList = children[c];
                    foreach (int ch in chList)
                    {
                        if (!controlledJoints.Contains(ch))
                        {
                            additionalCtl.Add(ch);
                        }
                    }
                }
                foreach (int c in additionalCtl)
                {
                    if (!controlledJoints.Contains(c))
                    {
                        controlledJoints.Add(c);
                    }
                }
            } while (additionalCtl.Count > 0);
            */

            for (int i = 0; i < skeleton.joints.Length; i++)
            {
                Matrix M = Matrix.CreateFromQuaternion(skeleton.joints[i].q);
                M.Translation = boneTransforms[i].Translation;
                boneTransforms[i] = M;
                //skeleton.boneTransforms.CopyTo(boneTransforms, 0);
            }
        }


        /// <summary>
        /// Helper used by the Update method to refresh the WorldTransforms data.
        /// </summary>
        public void UpdateWorldTransforms()
        {
            // Root bone.
            worldTransforms[0] = boneTransforms[0] * rootTransform;

            // Child bones.
            for (int bone = 1; bone < worldTransforms.Length; bone++)
            {
                int parentBone = skinningDataValue.SkeletonHierarchy[bone];

                worldTransforms[bone] = boneTransforms[bone] *
                                             worldTransforms[parentBone];
            }
        }


        /// <summary>
        /// Helper used by the Update method to refresh the SkinTransforms data.
        /// </summary>
        public void UpdateSkinTransforms()
        {
            for (int bone = 0; bone < skinTransforms.Length; bone++)
            {
                skinTransforms[bone] = skinningDataValue.InverseBindPose[bone] *
                                            worldTransforms[bone];
            }
        }


        /// <summary>
        /// Gets the current bone transform matrices, relative to their parent bones.
        /// </summary>
        public Matrix[] GetBoneTransforms()
        {
            return boneTransforms;
        }


        /// <summary>
        /// Gets the current bone transform matrices, in absolute format.
        /// </summary>
        public Matrix[] GetWorldTransforms()
        {
            return worldTransforms;
        }


        /// <summary>
        /// Gets the current bone transform matrices,
        /// relative to the skinning bind pose.
        /// </summary>
        public Matrix[] GetSkinTransforms()
        {
            return skinTransforms;
        }


        /// <summary>
        /// Gets the current play position.
        /// </summary>
        public TimeSpan CurrentTime
        {
            get { return currentTimeValue; }
        }


    }


    /// <summary>
    /// Basic Skeleton Class for represnting a joint hierarchy
    /// Stores joints and parent info.  Also stores joint naming (string->idx).
    /// </summary>
    public class Skeleton
    {
        

        public int[] parents; //parent info        
        public Joint[] joints;
        public Dictionary<string, int> boneNames;
        public Vector3 axis; //DEBUG TODO: remove this

        public Skeleton(Matrix[] boneTransforms, int[] parents, IDictionary<string, int> boneNames)
        {
            joints = new Joint[boneTransforms.Length];
            this.parents = new int[parents.Length];

            this.boneNames = new Dictionary<string, int>();

            parents.CopyTo(this.parents, 0);

            Dictionary<string, int>.KeyCollection keys = ((Dictionary<string, int>)boneNames).Keys;

            foreach (string k in keys)
            {
                //Console.WriteLine(k);
                int idx;
                idx = boneNames[k];
                this.boneNames.Add(k, idx);
            }
            Quaternion q = new Quaternion();
            Vector3 s = new Vector3();
            Vector3 t = new Vector3();

            for (int i = 0; i < joints.Length; i++)
            {
                boneTransforms[i].Decompose(out s, out q, out t);                

                q.Normalize();

                Matrix DBG = Matrix.CreateFromQuaternion(q);

                Vector3 v = new Vector3();
                TMath.Quaternion2Eulers(ref q, out v);
                v *= 180.0f/MathHelper.Pi;
                //Console.WriteLine("Initial rotations"); Console.WriteLine(v);

                joints[i] = new Joint(q, t);
                
            }
        }

        //concatenate joint transformations
        public void GetJointRotation(int idx, out Quaternion q)
        {
            q = new Quaternion();
            q = Quaternion.Identity;
            int ndx = idx;
            Stack<Quaternion> rotations = new Stack<Quaternion>();
            //while not at root
            while (ndx > -1)
            {
                rotations.Push(joints[ndx].q);
                ndx = parents[ndx];
            }

            while (rotations.Count > 0)
            {
#warning this may be a glitch
                q = Quaternion.Concatenate(q,Quaternion.Inverse(rotations.Pop()));
            }
            q.Normalize();
        }

        ///<summary>
        //chain together a joint and comptue the position of 
        //an end effector
        //TODO: speed this up
        /// </summary>
        /// <param name="idx"></param>
        /// <param name="v"></param>
        public void GetEffectorPosition(int idx, out Vector3 v)
        {
            v = new Vector3(0.0f, 0, 0);

            int ndx = idx;
            Stack<Joint> jointChain = new Stack<Joint>();

            //while not at root
            while (ndx > -1)
            {
                jointChain.Push(joints[ndx]);
                ndx = parents[ndx];
            }

            Joint j = jointChain.Pop();
            v = j.t;
            Quaternion prevRot = j.q;

            while (jointChain.Count > 0)
            {
                j = jointChain.Pop();
                v += Vector3.Transform(j.t, prevRot);
                prevRot *= j.q;
            }

        }

        /// <summary>
        /// TODO: Speed this up (see GetEffectorPosition
        /// </summary>
        /// <param name="idx"></param>
        /// <param name="aim"></param>
        public void GetEffectorAim(int effectorId, int jointId, out Vector3 aim)
        {
            Vector3 parent = new Vector3();
            Vector3 effector = new Vector3();

            //TODO: combine these into something more efficient
            GetEffectorPosition(jointId, out parent);

            GetEffectorPosition(effectorId, out effector);

            aim = effector - parent;
            aim.Normalize();
        }

        /// <summary>
        /// TODO: Speed this up (see GetEffectorPosition
        /// </summary>
        /// <param name="idx"></param>
        /// <param name="aim"></param>
        public void GetEffectorAim(int effectorId, out Vector3 aim)
        {
            int jointId = parents[effectorId];

            GetEffectorAim(effectorId, jointId, out aim);
        }

    }




    /// <summary>
    /// Basic joint class for our IK system
    /// 
    /// </summary>    
    public class Joint
    {
        public Joint(Quaternion q, Vector3 t)
        {
            this.q = q;
            this.t = t;
        }
        public Quaternion q = new Quaternion();
        public Vector3 t = new Vector3();
    }


}
